package main

import (
	"fmt"
	"sync"
	"time"
)

func main() {
	tokenBucket := NewTokenBucket(300, 500)
	succent := 0
	faild := 0
	for i := 0; i < 10000; i++ {
		if tokenBucket.TakeToken(1) {
			fmt.Println("Request processed.")
			succent += 1
		} else {
			fmt.Println("Request dropped.")
			faild += 1
		}

		time.Sleep(time.Millisecond)
	}

	fmt.Printf("成功%d次,失败%d 次", succent, faild)
}

type TokenBucket struct {
	rate           float64    //令牌放入速率 单位令牌/秒
	capacity       float64    //令牌桶容量
	tokens         float64    //当前令牌个数
	lastUpdateTime time.Time  //最近一次更新时间
	lock           sync.Mutex //保证并发安全
}

func NewTokenBucket(rate, capacity float64) *TokenBucket {
	return &TokenBucket{
		rate:           rate,
		capacity:       capacity,
		tokens:         capacity,
		lastUpdateTime: time.Now(),
	}
}

func (tb *TokenBucket) getToken() float64 {
	//tb.lock.Lock()
	//defer tb.lock.Unlock()

	now := time.Now()
	elapsed := now.Sub(tb.lastUpdateTime).Seconds()
	tb.lastUpdateTime = time.Now()

	tb.tokens += elapsed * tb.rate

	if tb.tokens > tb.capacity {
		tb.tokens = tb.capacity
	}

	return tb.tokens
}

func (tb *TokenBucket) TakeToken(amount float64) bool {
	tb.lock.Lock()
	defer tb.lock.Unlock()

	//这个时候增加 token 的数量
	tokens := tb.getToken()
	fmt.Println("当前令牌桶中的令牌数量：", tokens)
	if tokens >= amount {
		tb.tokens -= amount
		return true
	}

	return false
}
